package com.tamude.gpt;

import com.alibaba.excel.EasyExcel;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import constant.MessageRoleEnum;
import entity.ChatMessage;
import okhttp3.*;

import java.io.IOException;
import java.io.InputStream;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.*;
import java.util.concurrent.TimeUnit;

public class EmbeddingQASystem {
    private static final String CLAUDE_API_KEY = "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";
    private static final String OPENAI_API_KEY = "sk-proj-ZtWf4lOx7ipk0Xo0pwdJT3BlbkFJLuJQsVPK8l5gzj204rZv";
    private static final String CLAUDE_ENDPOINT = "https://api.anthropic.com/v1/messages";
    private static final String OPENAI_EMBEDDING_ENDPOINT = "https://api.openai.com/v1/embeddings";

    private final OkHttpClient client;
    private final Map<List<Double>, String> embeddingToText;

    public EmbeddingQASystem() {
        this.client = new OkHttpClient
                .Builder()
                .proxy(new Proxy(Proxy.Type.HTTP, new InetSocketAddress("127.0.0.1", 7890)))
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
        this.embeddingToText = new HashMap<>();
    }

    public List<Double> getEmbedding(String text) throws IOException {
        JSONObject requestBody = new JSONObject();
        requestBody.put("model", "text-embedding-3-small");
        requestBody.put("input", text);

        Request request = new Request.Builder()
                .url(OPENAI_EMBEDDING_ENDPOINT)
                .addHeader("Content-Type", "application/json")
                .addHeader("Authorization", "Bearer " + OPENAI_API_KEY)
                .post(RequestBody.create(MediaType.parse("application/json"), requestBody.toJSONString()))
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) throw new IOException("Unexpected code " + response);

            String responseBody = response.body().string();
            JSONObject jsonResponse = JSON.parseObject(responseBody);
            JSONArray embeddingArray = jsonResponse.getJSONArray("data").getJSONObject(0).getJSONArray("embedding");

            List<Double> embedding = new ArrayList<>();
            for (int i = 0; i < embeddingArray.size(); i++) {
                embedding.add(embeddingArray.getDoubleValue(i));
            }
            return embedding;
        }
    }

    public String generateAnswer(String prompt) throws IOException {
        JSONObject requestBody = new JSONObject();
        requestBody.put("model", "claude-3-sonnet-20240229");
        ChatMessage message = new ChatMessage(MessageRoleEnum.USER.getName(),prompt);
        requestBody.put("messages", Collections.singletonList(message));

        Request request = new Request.Builder()
                .url(CLAUDE_ENDPOINT)
                .addHeader("Content-Type", "application/json")
                .addHeader("x-api-key", CLAUDE_API_KEY)
                .post(RequestBody.create(MediaType.parse("application/json"), requestBody.toJSONString()))
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) throw new IOException("Unexpected code " + response);

            String responseBody = response.body().string();
            JSONObject jsonResponse = JSON.parseObject(responseBody);
            return jsonResponse.getString("completion");
        }
    }

    public void processAndStoreText(String text) throws IOException {
        List<Double> embedding = getEmbedding(text);
        embeddingToText.put(embedding, text);
    }

    public String findMostSimilarText(String query) throws IOException {
        List<Double> queryEmbedding = getEmbedding(query);
        double maxSimilarity = Double.NEGATIVE_INFINITY;
        String mostSimilarText = "";

        for (Map.Entry<List<Double>, String> entry : embeddingToText.entrySet()) {
            double similarity = cosineSimilarity(queryEmbedding, entry.getKey());
            if (similarity > maxSimilarity) {
                maxSimilarity = similarity;
                mostSimilarText = entry.getValue();
            }
        }

        return mostSimilarText;
    }

    private double cosineSimilarity(List<Double> v1, List<Double> v2) {
        double dotProduct = 0.0;
        double norm1 = 0.0;
        double norm2 = 0.0;
        for (int i = 0; i < v1.size(); i++) {
            dotProduct += v1.get(i) * v2.get(i);
            norm1 += Math.pow(v1.get(i), 2);
            norm2 += Math.pow(v2.get(i), 2);
        }
        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }

    public String answerQuestion(String question) throws IOException {
        String relevantText = findMostSimilarText(question);
        String prompt = "Human: Based on the following information:\n\n" + relevantText + "\n\nPlease answer this question: " + question + "\n\nAssistant: ";
        return generateAnswer(prompt);
    }

    public String readExcel(){
        InputStream inputStream = this.getClass().getClassLoader().getResourceAsStream("embedding测试数据.xlsx");
        List<TestEntity> list = EasyExcel.read(inputStream)
                .head(TestEntity.class)
                // 设置sheet,默认读取第一个
                .sheet()
                // 设置标题所在行数
                .headRowNumber(0)
                .doReadSync();
        StringBuilder sb = new StringBuilder();
        list.forEach(e -> {
            sb.append(e.getUser()).append(" | ").append(e.getCommodity()).append(" | ").append(e.getBuyDate()).append(" | ").append(e.getStars()).append(" | ").append(e.getEvaluation()).append(" | ").append(e.getEvaluationDate()).append("\n");
        });
        return sb.toString();
    }

    public static void main(String[] args) {
        EmbeddingQASystem qaSystem = new EmbeddingQASystem();
        try {

            // 处理和存储一些示例文本
            qaSystem.processAndStoreText(qaSystem.readExcel());

            // 尝试回答一个问题
            String question = "游戏宅在什么时间购买了什么商品？他对该商品的评价如何呢？";
            String answer = qaSystem.answerQuestion(question);
            System.out.println("Question: " + question);
            System.out.println("Answer: " + answer);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}